import scipy.io as spio
import scipy.sparse as sp
import numpy
import collections
import operator
import sys
import getopt
import os
import random_walk

"""
This script implements the random walk on a multi-plane network

J. He
06-28-2011
"""

#planes names
planes = []  
#teleport probability between planes
beta = [] 
#self-transition
alpha = 0
#number of steps of walk, default 5 
steps = 5
#walk types: forward or backward
walktype = 'backward' 
#post-processing
merge_nodes = 1

inputdir = ''
outputdir = ''


def get_plane(planename, inputdir):
	f = open('%s/%s.nodes'%(inputdir, planename))
	nodes = [n.strip().split(' ', 1) for n in f.readlines()]
	nodes = [(int(n[0]), n[1]) for n in nodes]
	f.close()
	tmp = [n[1] for n in nodes]
	if not len(set(tmp)) == len(tmp):
		print "WARNING: there exist duplicated node names: they are treated as different nodes in the local graph"
	f = '%s/%s.trans'%(inputdir, planename)
	cols, rows, p = numpy.loadtxt(f, dtype={'names': ('id1', 'id2', 'score'), "formats": ('i', 'i', 'f')}, unpack=True)
	mat = sp.coo_matrix((p, (rows, cols)), shape=(len(nodes), len(nodes))).tolil() 
	print "Renormalizing the row probabilities, just in case"
	normalizer = numpy.tile(mat.sum(1),(1, len(nodes)))
	mat = sp.coo_matrix(mat.multiply(1/normalizer))
	return nodes, mat 


"""
p(d_k|d_l) = \sum_{i,j} p(d_k|q_i)p_{t|0}(q_i|q_j)p(q_j|d_l)

construct the multi-plane network with teleport transition
in a network, assume node 1,2,3 are in plane 1, 4, 5 are in plane 2, and 2, 4 are the same
alpha as self transition, beta1(2) is the probability of teleport to plane 1 (2)

The 1-step transition matrix contains
	1			2			3			4		5			
[ alpha	           	(1-alpha)p(2|1) 	(1-alpha)p(3|1) 		0		0		
 (1-alpha)beta1*p(1|2)  alpha*beta1		(1-alpha)beta1*p(3|2)		alpha*beta2	(1-alpha)beta2*p(5|4)	
 (1-alpha)p(1|3)	(1-alpha)p(2|3)		alpha				0		0				 	
 (1-alpha)beta1*p(1|2)	alpha*beta1		(1-alpha)beta1*p(3|2)		alpha*beta2	(1-alpha)beta2*p(5|4)	
 0			0			0				(1-alpha)p(4|5)	alpha				
]

In summary, the transition probability from node i to node j:

P(j^a|i^b) = 
1) 0  (if i is not a teleport point and a != b)

2) alpha*beta_a for self-transition, (1-alpha)beta_a*p(j|k) for non-self-transition 
(if i is a teleport point) - where k is the telport target point 

3) alpha for self-transition, (1-alpha)p_ij for non-self-transition 
(if i is not a teleport point and a == b)

"""

def construct_graph_with_teleport(planes, alpha, beta):
	#prepare the transition matrix
	nodes = []
	nodes_plane = []
	trans = [i for i in range(len(planes))] 
	for i in range(len(planes)):
		#empty plane
		plane_nodes = planes[i][0]
		if len(plane_nodes)==0:
			continue
		x = len(nodes)
		#assing new node_id to each node in the network
		#each node contains a tuple (network_id, node_name, localgraph_id)
		nodes[x:] = [(n+x, plane_nodes[n][1], plane_nodes[n][0]) for n in range(len(plane_nodes))]
		nodes_plane[x:] = [(n+x, i) for n in range(len(plane_nodes))]
		trans[i] = (i, planes[i][1].todok())
	nodes_plane = dict(nodes_plane)
	trans = dict(trans)
	if len(nodes) == 0:
		return [],[]
	#find teleport points
	nodes.sort(key = operator.itemgetter(1))	
	tel = []
	current = nodes[0][1]
	tmp = [nodes[0][0]] 
	for n in nodes[1:]:
		if n[1]!= current:
			if len(tmp)>1:
				tel.append((current, tmp))	
			current = n[1]	  
			tmp = []
		tmp.append(n[0])
	if len(tmp)>1:
		tel.append((current, tmp))
	tel = dict(tel)
	nodes.sort(key=operator.itemgetter(0))
	print "Teleport points:", tel

	#create the matrix i in the row, j in the columns, fill in p(j|i)
	data = collections.deque()
	ind = 0
	for i in range(len(nodes)):
		plane_i = nodes_plane[nodes[i][0]] 
		#check and get the target of a teleport point 
		tel_target = tel.get(nodes[i][1], [])
		tel_target_planes = [nodes_plane[t] for t in tel_target]
		#renormalize the plane probability with the ones that are actually considered
		normalizer = sum([beta[t] for t in tel_target_planes])
		#we need to normalize the non-self transitions summing up to 1-alpha 
		sum_nonself = 0
		nonself = []
		for j in range(len(nodes)):
			plane_j = nodes_plane[nodes[j][0]]
			p_ji = 0
			if tel_target!=[]:
				#i is a telepoint, every entry needs to have its teleport probability
				b = beta[plane_j]/normalizer
				#only if j is in the target plane, otherwise it's not connected
				if plane_j in tel_target_planes:
					#get the target point
					k = tel_target[tel_target_planes.index(plane_j)]
					#trans[x] is a dok matrix, get((row, col)) gets a value
					p_jk = trans[plane_j].get((nodes[k][2], nodes[j][2]), 0)
					if k == j:
						p_ji = b*alpha
					else:	
						p_ji = b*p_jk*(1-alpha)	
				#		#non-self transitions are normalized to 1-alpha later
				#		#also if alpha = 0, self-transition + non-self-transiino =1
						if p_ji > 0:
							nonself.append(ind)
							sum_nonself += p_ji	
			else:
				#i is not a teleport point, only nodes on the same plane can be connected
				if plane_i == plane_j:
					p_ji = trans[plane_i].get((nodes[i][2], nodes[j][2]), 0)
					if j == i:
						p_ji = alpha
					else:
						p_ji = p_ji*(1-alpha)
						#otherwise nonself includes the original self-trans given in the graphs
						if p_ji > 0:
							nonself.append(ind)
							sum_nonself += p_ji
			if p_ji > 0:
				data.append((i, j, p_ji))
				ind += 1
		#normalize the non-self transition to 1-alpha
		for x in nonself:
			score = data[x][2]
			if sum_nonself >0:
				score = data[x][2]/sum_nonself * (1-alpha)
			data[x] = (data[x][0], data[x][1], score)

	ind_row = [d[0] for d in data] 
	ind_col = [d[1] for d in data]
	D = [d[2] for d in data]
	matrix = sp.coo_matrix((D, (ind_row, ind_col)), shape=(len(nodes), len(nodes))).tocsr()				
	return matrix, nodes, tel


def load_config(inputfile):
	global planes
	global beta 
	global outputdir
	global inputdir
	global alpha
	global steps
	global walktype
	global merge_nodes

	f = open(inputfile)	
	params = []	 
        for c in f:
                if c.strip().startswith('#'):
                        continue
                strs = c.split('#')[0].split('=')
                if len(strs)==2:
                        params.append([s.strip() for s in strs])
        f.close()	
	params = dict(params)
	keys = params.keys()
	if 'planes' in keys:
		planes = params['planes'].split(';')
	else:
		print "Error: tables of planes needs to be provided, separated by ';'"
	#default
	beta = [float(1)/float(len(planes)) for p in planes]
	if 'beta' in keys:
		teleport = [float(t) for t in params['beta'].split(';')]
		if not len(teleport) == len(planes):
			print "Error: length of beta should equal to the length of planes"
			sys.exit(2) 
		if teleport == []:
			teleport = [1 for p in planes]
		sum_tmp = sum(teleport)
		beta = [float(t)/float(sum_tmp) for t in teleport]
	if 'outputdir' in keys:
		outputdir = params['outputdir']
	else:
		print 'outputdir should be provided'
		sys.exit(2)
        if 'inputdir' in keys:
                inputdir = params['inputdir']
        else:
                print 'inputdir should be provided'
		sys.exit(2)
	if 'alpha' in keys:
		alpha = float(params['alpha'])
	if 'steps' in keys:
		steps = int(params['steps'])
	if 'walktype' in keys:
		walktype = params['walktype']
	if 'merge_nodes' in keys:
		merge_nodes = int(params['merge_nodes'])

def usage():
        print "usage: python multiplane_walk.py -i configfile"
        print "-h               print this message"

def main(argv=None):
        try:
                opts, args = getopt.getopt(sys.argv[1:], "hi:")
        except getopt.GetoptError, err:
                # print help information and exit:
                print str(err) # will print something like "option -a not recognized"
                usage()
                sys.exit(2)
        configfile = ''
        if len(opts)==0:
                usage()
                sys.exit()

        for o, a in opts:
                if o == "-h":
                        usage()
                        sys.exit()
                elif o == '-i':
                        configfile = a
                else:
                        assert False, "unhandled option"

        if not os.path.exists(configfile):
                print "config file does not exist"
                sys.exit()
	load_config(configfile)	
		
	output = '%s/planes%s_alpha%s_step%s'%(outputdir, '_'.join(planes),alpha, steps)
	output_nodes = '%s.nodes'%output
	output_trans = '%s.trans'%output
	if not os.path.exists(outputdir):
		os.mkdir(outputdir)

	all_planes = []	
	for plane in planes:	
		print "Load graph %s."%(plane)
		p = get_plane(plane, inputdir)
		all_planes.append(p)
	print 'Constructing planed graph...'
	mat, nodes, tel = construct_graph_with_teleport(all_planes, alpha, beta)
	print 'Random walk with settings:'
	print 'alpha:', alpha, 'steps:', steps, 'walktype:', walktype, 'beta:', beta
	if len(nodes)>0: 
	#the starting/ending matrix is an Identity matrix as we need all nodes participate
		start = sp.identity(mat.shape[0])
		rw = random_walk.random_walk()
		rw.load_transition_matrix(mat)
		if walktype == 'forward':
			res = rw.backward_walk(start, steps)	
		else:
			res = rw.backward_walk(start, steps)	
	res = res.todok()
	scores = collections.deque()
	#for teleport points, the same string occurs in different planes.
	#The probability of transit from a particular string to another can be merged.
	if merge_nodes == 1:
		#n as row
		for n in nodes:
			if n[1] in tel:
				#row is telpoint, all entries are repeated k times, k=#connected telepoints
				#only perseve the first one, ignore others
				if not n[0] == tel[n[1]][0]:
					continue
			#m as col
			for m in nodes:
				score = 0
				if m[1] in tel:
					# if col is a telpoint, sum the scores from all occurrence, only store it once
					if m[0] == tel[m[1]][0]:
						for t in tel[m[1]]:
							score += res.get((n[0], t), 0) 	
					else:
						continue	
				else:
					score = res.get((n[0], m[0]), 0)
				scores.append((n[1], m[1], score))		 
	else:
		print "no merge"
		for n in nodes:
			for m in nodes:
				score = res.get((n[0], m[0]), 0)
				scores.append((n[1], m[1], score))

	nodes_unique = sorted(list(set([n[1] for n in nodes])))
	unique_ids = dict([(nodes_unique[i], i) for i in range(len(nodes_unique))])

	out_nodes = open(output_nodes, 'w')
	for n in nodes_unique:
		out_nodes.write('%s %s\n'%(unique_ids[n], n))	
	out_nodes.close()

	out_trans = open(output_trans, 'w')
	scores = list(scores)
	scores.sort(key=operator.itemgetter(0))
	for s in scores:
		row = unique_ids.get(s[0])
		col = unique_ids.get(s[1]) 
		out_trans.write('%s %s %s\n'%(col, row, s[2]))	
	out_trans.close()
	print 'Saving result to:\n', output_nodes, output_trans



if __name__ == "__main__":
    sys.exit(main())








